import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 设置随机种子以确保结果可复现
torch.manual_seed(42)

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
])

# 加载MNIST数据集
train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform)

# 创建数据加载器
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=1000, shuffle=False)

# 定义CNN模型


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(1600, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = self.conv2_drop(x)
        x = x.view(-1, 1600)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return torch.log_softmax(x, dim=1)


# 实例化模型
model = CNN().to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型


def train(epochs=5):
    model.train()
    train_losses = []
    for epoch in range(epochs):
        running_loss = 0.0
        for i, (images, labels) in enumerate(train_loader):
            images, labels = images.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            if i % 100 == 99:
                print(f'[{epoch + 1}, {i + 1}] 损失: {running_loss / 100:.3f}')
                train_losses.append(running_loss / 100)
                running_loss = 0.0

    print('训练完成!')
    return train_losses

# 测试模型


def test():
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            test_loss += criterion(outputs, labels).item()
            pred = outputs.argmax(dim=1, keepdim=True)
            correct += pred.eq(labels.view_as(pred)).sum().item()

    test_loss /= len(test_loader)
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'测试集平均损失: {test_loss:.4f}, 准确率: {accuracy:.2f}%')
    return accuracy

# 可视化一些测试结果


def visualize_results():
    model.eval()
    with torch.no_grad():
        # 获取一批测试数据
        dataiter = iter(test_loader)
        images, labels = next(dataiter)
        images, labels = images.to(device), labels.to(device)

        # 获取预测结果
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        # 显示前10个图像及其预测结果
        fig = plt.figure(figsize=(12, 4))
        for i in range(10):
            ax = fig.add_subplot(1, 10, i + 1)
            ax.imshow(images[i].cpu().squeeze(), cmap='gray')
            ax.set_title(f'预测: {predicted[i].item()}\n实际: {labels[i].item()}')
            ax.axis('off')
        plt.tight_layout()
        plt.savefig('mnist_results.png')
        plt.close()

# 主函数


def main():
    # 训练模型
    print("开始训练模型...")
    train_losses = train(epochs=5)

    # 测试模型
    print("\n开始测试模型...")
    accuracy = test()
    print(f"最终测试准确率: {accuracy:.2f}%")

    # 可视化结果
    print("\n生成可视化结果...")
    visualize_results()

    # 保存模型
    torch.save(model.state_dict(), 'mnist_cnn_model.pth')
    print("模型已保存到 mnist_cnn_model.pth")

    # 绘制训练损失曲线
    plt.figure(figsize=(10, 5))
    plt.plot(train_losses)
    plt.title('训练损失曲线')
    plt.xlabel('迭代次数 (x100)')
    plt.ylabel('损失')
    plt.savefig('training_loss.png')
    plt.close()
    print("训练损失曲线已保存到 training_loss.png")


if __name__ == "__main__":
    main()
